/*
 * Copyright (c) IPADS@SJTU 2021. All rights reserved.
 * secGear is licensed under the Mulan PSL v2.
 * You can use this software according to the terms and conditions of the Mulan PSL v2.
 * You may obtain a copy of Mulan PSL v2 at:
 *     http://license.coscl.org.cn/MulanPSL2
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY OR FIT FOR A PARTICULAR
 * PURPOSE.
 * See the Mulan PSL v2 for more details.
 */

#include <stdio.h>
#include <stddef.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <limits.h>

#include "enclave.h"
#include "enclave_internal.h"
#include "enclave_log.h"
#include "msg_structure.h"
#include "client_run.h"
#include "qingtian_enclave.h"

#define DBUG_MOCK

extern list_ops_management g_list_ops;

#define BUF_PRA_SIZE        (64)
#define CMD_BUF_PRA_SIZE    (128)
#define CMD_BUF_MAX         (128 + PATH_MAX)

#define CMD_BUF_RESULT_MAX      (1024 * 5)
#define CMD_BUF_RESULT_LINE_MAX (128)

#define CID_MIN             (4)
// return cmd length
static int get_start_cmdline(char *cmd_buf, size_t cmd_buf_size, const qingtian_startup_t* pra,
                             const char *path, bool debug_mode)
{
    char *cmd = "qt enclave start";
    char buf[BUF_PRA_SIZE];
    if (cmd_buf == NULL || cmd_buf_size == 0) {
        print_error_term("command buf is NULL\n");
        return -1;
    }
    if (path == NULL) {
        print_error_term("command parameter path is NULL\n");
        return -1;
    }
    if (cmd_buf_size - 1 < strlen(cmd) + strlen(path) + CMD_BUF_PRA_SIZE) {
        return -1;
    }
    (void)sprintf(cmd_buf, "%s --eif %s", cmd, path);
    if (debug_mode) {
        strcat(cmd_buf, " --debug-mode");
    }
    if (pra == NULL) {
        goto end;
    }
    if (pra->cpus != 0) {
        sprintf(buf, " --cpus %lu", (long unsigned int)pra->cpus);
        strcat(cmd_buf, buf);
    }
    if (pra->mem_mb != 0) {
        sprintf(buf, " --mem %lu", (long unsigned int)pra->mem_mb);
        strcat(cmd_buf, buf);
    }
    if (pra->enclave_cid >= CID_MIN) {
        sprintf(buf, " --cid %lu", (long unsigned int)pra->enclave_cid);
        strcat(cmd_buf, buf);
    } else {
        print_error_term("enclave id must not less %d\n", CID_MIN);
        return -1;
    }
end:
    return strlen(cmd_buf);
}

static int check_eif(const char* path)
{
    if (access(path, F_OK) != 0) {
        print_error_term("%s can not access\n", path);
        return -1;
    }
    // file extern must be "eif"
    char *ext = strrchr(path, '.');
    if (ext == NULL || strcmp(ext, ".eif") != 0) {
        print_error_term("%s must be end of .eif\n", path);
        return -1;
    }
    return 0;
}

static int get_eif_realpath(char *resolved_path, const char *path)
{
    // TODO 参数合法性检查
    int ret = 0;
    if (path == NULL) {
        print_error_term("path is NULL\n");
        ret = -1;
        goto end;
    }
    if (realpath(path, resolved_path) == NULL) {
        print_error_term("%s realpath fail\n", path);
        ret = -1;
        goto end;
    }
    if (check_eif(resolved_path) != 0) {
        print_error_term("check %s fail\n", resolved_path);
        ret = -1;
        goto end;
    }
end:
    return ret;
}

// return length of command string
int qt_start_cmd_construct(char *command, const qingtian_startup_t *pra, const char *eif, uint32_t flags)
{
    int ret = 0;
    char *resolved_path = calloc(1, PATH_MAX);
    if (resolved_path == NULL) {
        ret = -1;
        goto end;
    }
    if (get_eif_realpath(resolved_path, eif) != 0) {
        ret = -1;
        goto end;
    }
    bool debug = false;
    if (flags & SECGEAR_DEBUG_FLAG) {
        debug = true;
    } else {
        debug = false;
    }
    ret = get_start_cmdline(command, CMD_BUF_MAX, pra, resolved_path, debug);
    if (ret < 0) {
        print_error_term("get start command fail\n");
        ret = -1;
        goto end;
    }
end:
    if (resolved_path != NULL) {
        free(resolved_path);
    }
    return ret;
}

static int get_id(char *str, unsigned int *cid, unsigned int *id)
{
    char *target = NULL;
    int ret = 0;
    target = strstr(str, "EnclaveCID");
    if (target == NULL) {
        return -1;
    }
    ret = sscanf(target, "%*[^:]:%u", cid);
    if (ret != 1) {
        return -1;
    }

    target = strstr(str, "EnclaveID");
    if (target == NULL) {
        return -1;
    }
    ret = sscanf(target, "%*[^:]:%u", id);
    if (ret != 1) {
        return -1;
    }
    return 0;
}

static int get_match_id(char *str, unsigned int cid, unsigned int *id)
{
    char *cur = NULL;
    char *next = NULL;
    unsigned int tmp_cid;
    unsigned int tmp_id;
    const char *delimeter = "}";
    cur = strtok_r(str, delimeter, &next);
    while (cur != NULL) {
        if (get_id(cur, &tmp_cid, &tmp_id) != 0) {
            cur = strtok_r(NULL, delimeter, &next);
        }
        if (cid == tmp_cid) {
            *id = tmp_id;
            return 0;
        }
        cur = strtok_r(NULL, delimeter, &next);
    }
    return -1;
}

static int qt_query_id(unsigned int cid, unsigned int *id)
{
    int ret = 0;
    FILE *fp = NULL;
    const char *cmd = "qt enclave query";
    char read_buf[CMD_BUF_RESULT_LINE_MAX];
    char *buf = calloc(1, CMD_BUF_RESULT_MAX);
    if (buf == NULL) {
        ret = -1;
        goto end;
    }
    fp = popen(cmd, "r");
    if (fp == NULL) {
        ret = -1;
        goto end;
    }
    while (fgets(read_buf, CMD_BUF_RESULT_LINE_MAX, fp) != NULL) {
        strcat(buf, read_buf);
    }
    print_debug("qt query: %s \n", buf);
    if (get_match_id(buf, cid, id) != 0) {
        print_debug("cid = %u, get id fail\n", cid);
        ret = -1;
        goto end;
    }
    print_debug("cid = %u, get id = %u\n", cid, *id);
end:
    if (fp != NULL) {
        fclose(fp);
    }
    if (buf != NULL) {
        free(buf);
    }
    return ret;
}

static int qt_start(char *command, unsigned int cid, uint32_t *id, int retry)
{
    FILE *fp = NULL;
    int ret = 0;
    if (command == NULL || id == NULL) {
        ret = -1;
        goto end;
    }
    print_debug("exec cmd: %s\n", command);
    fp = popen(command, "r");
    if (fp == NULL) {
        print_error_term("command execute failed\n");
        ret = -1;
        goto end;
    } else {
        print_debug("get enclave id, total retry %d\n", retry);
        while (retry-- > 0) {
            print_debug("try %d\n", retry + 1);
            if (qt_query_id(cid, id) != 0) {
                sleep(1);
                continue;
            } else {
                break;
            }
        }
        if (retry <= 0) {
            ret = -1;
            print_debug("query id fail\n");
        } else {
            ret = 0;
            print_debug("qingtian enclave id  %lu\n", *id);
        }
    }
end:
    if (fp != NULL) {
        pclose(fp);
        fp = NULL;
    }
    return ret;
}

cc_enclave_result_t _qingtian_create(cc_enclave_t *enclave, const enclave_features_t *features,
                                     const uint32_t features_count)
{
    cc_enclave_result_t result_cc = CC_SUCCESS;
    char *command = NULL;
    if (enclave == NULL) {
        print_error_term("Context parameter is NULL\n");
        return CC_ERROR_BAD_PARAMETERS;
    }
    qingtian_startup_t *startup_pra = NULL;
    for (uint32_t index = 0; index < features_count; index++) {
        if (features[index].setting_type & QINGTIAN_STARTUP_FEATURES) {
            startup_pra = features[index].feature_desc;
            break;
        }
    }
    if (startup_pra == NULL) {
        print_error_term("enclave startup parameter is NULL\n");
        return CC_ERROR_BAD_PARAMETERS;
    }
    command = calloc(1, CMD_BUF_MAX);
    if (command == NULL) {
        print_error_term("malloc for start command is NULL\n");
        result_cc = CC_ERROR_OUT_OF_MEMORY;
        goto end;
    }
    if (qt_start_cmd_construct(command, startup_pra, enclave->path, enclave->flags) <= 0) {
        print_error_term("construct qt start command fail\n");
        result_cc = CC_ERROR_GENERIC;
        goto end;
    }
    print_debug("exec cmd: %s\n", command);
    uint32_t id;
    if (qt_start(command, (unsigned int)startup_pra->enclave_cid, &id, startup_pra->query_retry) != 0) {
        print_debug("qingtian enclave start command exec fail!\n");
        result_cc = CC_ERROR_GENERIC;
        goto end;
    }
    if (client_init((unsigned int)startup_pra->enclave_cid, startup_pra->port) != 0) {
        print_error_term("try connect to enclave fail\n");
        result_cc = CC_ERROR_GENERIC;
        goto end;
    }
    print_debug("qingtian enclave create successfully! \n");
    qingtian_private_data_t *priv_data = (qingtian_private_data_t *)malloc(sizeof(qingtian_private_data_t));
    priv_data->enclave_id = id;
    priv_data->startup = *startup_pra;
    enclave->private_data = (void *)priv_data;
    result_cc = CC_SUCCESS;
end:
    if (command != NULL) {
        free(command);
    }
    return result_cc;
}

static int qt_stop(uint32_t enclave_id)
{
    int ret = 0;
    FILE *fp = NULL;
    char *target = NULL;
    char *cmd_result = NULL;
    char *command = calloc(1, CMD_BUF_MAX);
    if (command == NULL) {
        ret = -1;
        goto end;
    }
    ret = sprintf(command, "qt enclave stop --enclave-id %lu", (long unsigned int)enclave_id);
    if (ret <= 0) {
        print_error_term("host: construct command fail\n");
        ret = -1;
        goto end;
    }
    cmd_result = calloc(1, CMD_BUF_MAX);
    if (cmd_result == NULL) {
        ret = -1;
        goto end;
    }
    print_debug("exec cmd: %s\n", command);
    fp = popen(command, "r");
    if (fp == NULL) {
        print_error_term("popen failed\n");
        ret = -1;
        goto end;
    }
    ret = -1;
    while (fgets(cmd_result, CMD_BUF_MAX, fp) != NULL) {
        target = strstr(cmd_result, "success");
        if (target == NULL) {
            continue;
        }
        ret = 0;
        break;
    }
end:
    if (fp != NULL) {
        pclose(fp);
    }
    if (command != NULL) {
        free(command);
    }
    if (cmd_result != NULL) {
        free(cmd_result);
    }
    return ret;
}

cc_enclave_result_t _qingtian_destroy(cc_enclave_t *context)
{
    cc_enclave_result_t result_cc;
    if (context == NULL) {
        print_error_term("qinttian destroy parameter error\n");
        result_cc = CC_ERROR_BAD_PARAMETERS;
        goto end;
    }
    client_fini();

    qingtian_private_data_t *priv_data = (qingtian_private_data_t *)context->private_data;
    if (priv_data == NULL) {
        result_cc = CC_ERROR_BAD_PARAMETERS;
        goto end;
    }
    if (qt_stop(priv_data->enclave_id) != 0) {
        result_cc = CC_ERROR_GENERIC;
        goto end;
    }
    if (priv_data != NULL) {
        free(priv_data);
        context->private_data = NULL;
    }
    print_debug("qingtian destroy success\n");
    result_cc = CC_SUCCESS;
end:
    return result_cc;
}

// cc_enclave_result_t handle_ocall(
//     struct PLenclave* penglai_enclave,
//     uint8_t* untrusted_mem_extent,
//     const void *ocall_table,
//     int* result)
// {
//     untrusted_mem_info_t* ocall_mem_info;
//     size_t ocall_buf_size = 0;
//     int ocall_table_size = 0;
//     int ocall_func_id = 0;
//     uint8_t* ocall_in_buf = NULL;
//     uint8_t* ocall_out_buf = NULL;

//     ocall_mem_info = (untrusted_mem_info_t*)untrusted_mem_extent;
//     ocall_buf_size = size_to_aligned_size(sizeof(untrusted_mem_info_t)) +
//             ocall_mem_info->in_buf_size + ocall_mem_info->out_buf_size;
//     if(ocall_buf_size > DEFAULT_UNTRUSTED_SIZE){
//         print_debug("[ERROR]: the size of ocall parameters is too big to transfer through untrusted memory\n");
//         return CC_FAIL;
//     }

//     ocall_table_size = ((ocall_enclave_table_t*)ocall_table)->num;
//     ocall_func_id = ocall_mem_info->fid;
//     if(ocall_func_id >= ocall_table_size){
//         print_debug("[ERROR] host: ocall function isn't exist!\n");
//         return CC_FAIL;
//     }

//     ocall_in_buf = untrusted_mem_extent +
//             size_to_aligned_size(sizeof(untrusted_mem_info_t));
//     ocall_out_buf = ocall_in_buf + ocall_mem_info->in_buf_size;
//     if(((ocall_enclave_table_t*)ocall_table)->ocalls[ocall_func_id](
//             ocall_in_buf,
//             ocall_mem_info->in_buf_size,
//             ocall_out_buf,
//             ocall_mem_info->out_buf_size) != CC_SUCCESS){
//         print_debug("[ERROR] host: ocall function return false!\n");
//         return CC_FAIL;
//     }
//     penglai_enclave->user_param.ocall_buf_size = ocall_buf_size;
//     penglai_enclave->user_param.resume_type = USER_PARAM_RESUME_FROM_CUSTOM_OCALL;
//     *result = PLenclave_resume(penglai_enclave);

//     return CC_SUCCESS;
// }

static int32_t msg_send_recv(uint8_t* send_buf, size_t send_buf_size, uint8_t* recv_buf, size_t recv_buf_size)
{
    size_t recv_len = recv_buf_size;
    int ret_val = 0;
    int ret;
    ret = client_handle(send_buf, send_buf_size, recv_buf, &recv_len, &ret_val);
    if (ret != 0 || ret_val != CC_SUCCESS) {
        return -1;
    }
    return recv_len;
}

cc_enclave_result_t cc_enclave_call_function(
    cc_enclave_t *enclave,
    uint32_t function_id,
    const void *input_buffer,
    size_t input_buffer_size,
    void *output_buffer,
    size_t output_buffer_size,
    void *ms,
    const void *ocall_table)
{
    (void)enclave;
    (void)ms;
    (void)ocall_table;
    uint8_t *send_buf = NULL;
    uint32_t send_buf_size = 0;
    uint8_t* recv_buf = NULL;
    uint32_t recv_buf_size = 0;
#if 1
    print_debug("ecall input function id %lu\n", (long unsigned int)function_id);
    print_debug("ecall input buffer size %zu\n", input_buffer_size);
    print_debug("ecall output buffer size %zu\n", output_buffer_size);
    print_log(SECGEAR_LOG_LEVEL_DEBUG, "ecall input data: ");
    for (size_t i = 0; i < input_buffer_size; i++) {
        print_log(SECGEAR_LOG_LEVEL_DEBUG, "%02X", *((uint8_t*)input_buffer + i));
    }
    print_log(SECGEAR_LOG_LEVEL_DEBUG, "\n");
#endif
    // pack message into buffer
    cc_enclave_result_t result_cc = CC_SUCCESS;
    send_buf_size = MSG_BUF_SIZE(input_buffer_size);
    send_buf = calloc(1, send_buf_size);
    if (send_buf == NULL) {
        result_cc = CC_ERROR_OUT_OF_MEMORY;
        goto end;
    }
    msg_t msg_send;
    msg_send.function_id = function_id;
    msg_send.buf = (uint8_t *)input_buffer;
    msg_send.buf_size = input_buffer_size;
    msg_send.out_buf_size = output_buffer_size;
    int32_t msg_len = pack_msg(send_buf, send_buf_size, &msg_send);
    if (msg_len < 0) {
        result_cc = CC_ERROR_SHORT_BUFFER;
        goto end;
    }
    print_debug("message buffer[%lu]: ", send_buf_size);
#if 1
    for (int32_t i = 0; i < msg_len; i++) {
        printf("%02X", send_buf[i]);
    }
    printf("\n");
#endif
    // send and wait receive
    recv_buf_size = MSG_BUF_SIZE(output_buffer_size);
    recv_buf = calloc(1, recv_buf_size);
    if (recv_buf == NULL) {
        result_cc = CC_ERROR_OUT_OF_MEMORY;
        goto end;
    }
    int recv_len = 0;
    print_debug("msg send and wait...\n");
    recv_len = msg_send_recv(send_buf, send_buf_size, recv_buf, recv_buf_size);
    if (recv_len < 0) {
        result_cc = CC_ERROR_GENERIC;
        goto end;
    }
#if 1
    print_debug("received raw data[%zu]: ", recv_len);
    for (int i = 0; i < recv_len; i++) {
        printf("%02X", recv_buf[i]);
    }
    printf("\n");
#endif

    // unpack msg from receive buffer
    msg_t msg_recv;
    msg_recv.buf = NULL;
    msg_recv.buf_size= 0;
    int ret = unpack_msg(&msg_recv, recv_buf, recv_len);
    if (ret == ERR_MSG_GENERIC) {
        result_cc = CC_ERROR_SHORT_BUFFER;
        goto end;
    } else if (ret == ERR_MSG_SIZE) {
        result_cc = CC_ERROR_BAD_FORMAT;
        goto end;
    }

    // fill raw data into output buffer
    if (output_buffer_size < msg_recv.buf_size) {
        result_cc = CC_ERROR_SHORT_BUFFER;
        goto end;
    }
    print_debug("fill data(length %lu) into output buffer(size %lu): ", msg_recv.buf_size, output_buffer_size);
#if 1
    for(size_t i = 0; i < msg_recv.buf_size; i++) {
        printf("%02X", msg_recv.buf[i]);
    }
    printf("\n");
#endif
    (void)memcpy(output_buffer, msg_recv.buf, msg_recv.buf_size);
end:
    if (send_buf) {
        free(send_buf);
    }
    if (recv_buf) {
        free(recv_buf);
    }
    if (msg_recv.buf) {
        free(msg_recv.buf);
    }
    return result_cc;
}

const struct cc_enclave_ops global_qingtian_ops = {
    .cc_create_enclave = _qingtian_create,
    .cc_destroy_enclave = _qingtian_destroy,
    .cc_ecall_enclave = cc_enclave_call_function,
};

struct cc_enclave_ops_desc global_qingtian_ops_name = {
    .name = "qingtian",
    .ops = &global_qingtian_ops,
    .type_version = QINGTIAN_ENCLAVE_TYPE_0,
    .count = 0,
};

struct list_ops_desc global_qingtian_ops_node = {
    .ops_desc = &global_qingtian_ops_name,
    .next = NULL,
};

#define OPS_NAME global_qingtian_ops_name
#define OPS_NODE global_qingtian_ops_node
#define OPS_STRU global_qingtian_ops

cc_enclave_result_t cc_tee_registered(cc_enclave_t *context, void *handle)
{
    size_t len = strlen(OPS_NAME.name);
    if (OPS_NAME.type_version != context->type || OPS_NODE.ops_desc != &OPS_NAME ||
        len >= MAX_ENGINE_NAME_LEN || OPS_NAME.ops != &OPS_STRU) {
        print_error_term("The struct cc_enclave_ops_desc initialization error\n");
        return CC_ERROR_BAD_PARAMETERS;
    }
    OPS_NAME.handle = handle;
    context->list_ops_node = &OPS_NODE;
    add_ops_list(&OPS_NODE);
    return  CC_SUCCESS;
}

cc_enclave_result_t cc_tee_unregistered(cc_enclave_t *context, enclave_type_version_t type_version)
{
    if (context == NULL || context->list_ops_node != &OPS_NODE || type_version != OPS_NAME.type_version) {
        print_error_term("Engine parameter error \n");
        return CC_FAIL;
    }
    remove_ops_list(&OPS_NODE);
    return  CC_SUCCESS;
}
